package jtg.generator;

import jtg.DataFlowPath;
import jtg.DirectedGraph;
import jtg.ExpectResult;
import jtg.PrimePath;
import jtg.graphics.SootCFG;
import jtg.solver.Z3Solver;
import jtg.visualizer.Visualizer;
import org.apache.commons.lang3.RandomStringUtils;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIfStmt;
import soot.jimple.internal.JReturnStmt;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class SimpleGenerator {

    private String clsPath;
    private String clsName;
    private String mtdName;
    public UnitGraph ug;
    private Body body;
    public String pathType;


    public SimpleGenerator(String className, String methodName) {
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new SimpleGenerator(defaultClsPath, className, methodName,"1");
    }

    public SimpleGenerator(String classPath, String className, String methodName,String pathType) {
        clsPath = classPath;
        clsName = className;
        mtdName = methodName;
        ug = SootCFG.getMethodCFG(clsPath, clsName, mtdName);
        body = SootCFG.getMethodBody(clsPath, clsName, mtdName);
        this.pathType=pathType;
//        System.out.println(ug.toString());
//        System.out.println(body.toString());
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }


    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<>();
//        System.out.println("local::::::"+body.getLocals());
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameter() {
        ArrayList<Local> paras = new ArrayList<>();
        for (Local para : body.getParameterLocals()) {
            paras.add(para);
        }
        return paras;
    }

    public List<String> generate() {

        List<Unit> path;
        ArrayList<String> testSet = new ArrayList<>();
        ArrayList<String> resultSet=new ArrayList<>();
        String pathConstraint;


        DirectedGraph directedGraph = new DirectedGraph(ug.size());
        directedGraph.generateGraph(ug);
        
        List<List<Unit>> finalPath = new ArrayList<>();
        if(pathType.equals("1")) {
            List<List<Unit>> simplePath = PrimePath.calSimplePath(directedGraph);
            List<List<Unit>> primePath = PrimePath.calPrimePath(simplePath);
            List<List<Unit>> completePath = PrimePath.calCompletePath(primePath, directedGraph);
            List<String> num = PrimePath.toNumber(completePath, directedGraph);
            for (String s : num) {
                System.out.println(s);
            }
            finalPath=new ArrayList<>(completePath);
        }
        else if(pathType.equals("2")) {
            DataFlowPath dataFlowPath = new DataFlowPath(directedGraph, ug);
            List<List<Unit>> w1 = dataFlowPath.calDuPath(dataFlowPath.df.d, dataFlowPath.df.u);
            List<List<Unit>> w2 = new ArrayList<>(w1);
            for (List<Unit> l : w1) {
                for (List<Unit> m : w1) {
                    if (!m.equals(l)) {
                        if (l.toString().substring(1, (l.toString()).length() - 1).contains(m.toString().substring(1, m.toString().length() - 1))) {
                            w2.remove(m);
                        }
                    }
                }
            }
            List<List<Unit>> completePath = PrimePath.calCompletePath(w2, directedGraph);
            for (String u : PrimePath.toNumber(completePath, directedGraph)) {
                System.out.println(u);
            }
            finalPath=new ArrayList<>(completePath);
        }

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        List<Local> locals=body.getParameterLocals();
        try {
            testSet = new ArrayList<>();
            resultSet=new ArrayList<>();
            int round=0;
            for (List<Unit> l : finalPath) {
                path = l;
                System.out.println("The path is: " + path.toString());
                pathConstraint = calPathConstraint(path);
                //如果路径约束为空字符串，表示路径约束为恒真

//                List<Local> stringLocal=new ArrayList<>();
//                for(Local local:locals){
//                    if("java.lang.String".equals(local.getType().toString())){
//                        stringLocal.add(local);
//                    }
//                }

//                System.out.println(s);

                System.out.println("The corresponding path constraint is: " + pathConstraint);
                //The corresponding path constraint is: !( i0 % 15 != 0 )
                if (!pathConstraint.isEmpty() && !pathConstraint.contains("\"")) {
                    testSet.add(solve(pathConstraint));
                }
                else if (pathConstraint.contains("\""))
                {
                    testSet.add(solveString(locals.get(0),pathConstraint,round));
                }

                List<Local> free=new ArrayList<>();
                for(Local local:locals) {
                    if (!pathConstraint.contains(local.getName())) {
                        free.add(local);
                    }
                }
//                System.out.println("free"+free);
                if(!free.isEmpty()){
                    if(testSet.isEmpty()){
                        testSet.add(randomTC(free));
                    } else {
                        testSet.set(round, testSet.get(round) + randomTC(free));
                    }
                }
                if(!testSet.isEmpty()) {
                    testSet = sortTestSet(testSet, locals);
                }
//                System.out.println("test"+testSet);
                System.out.println("--------------------------------------------------------");
                round=round+1;
            }
        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
        }
        if (!testSet.isEmpty()) {
            System.out.println();
            resultSet= (ArrayList<String>) calExpectedResult(testSet,locals);
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) Test Input:" + tc);
                System.out.println("      Expected Result:"+resultSet.get(count-1));
                count++;
            }
        }
        System.out.println("Junit code：");
        String junit = ExpectResult.calJUnit(testSet, resultSet,mtdName, clsName.split("\\.")[1]);
        System.out.println(junit);
        return testSet;
    }

    public String calPathConstraint(List<Unit> path) {
        /* path:
         * r1 := @this: cut.LogicStructure,
         * i0 := @parameter0: int,
         * $i1 = i0 % 15,
         * if $i1 != 0 goto $i2 = i0 % 5,
         * return "FizzBuzz"
         */
        List<Local> jVars = getJVars();
        //jvars:[$i1, $i2, $i3, $r0]
//        System.out.println("jvars:"+jVars.toString());


        String pathConstraint = "";
        String expectedResult = "";

        HashMap<String, String> assignList = new HashMap<>();
        ArrayList<String> stepConditionsWithJimpleVars = new ArrayList<>();
        ArrayList<String> stepConditions = new ArrayList<>();

        for (Unit stmt : path) {

            if (stmt instanceof JAssignStmt) {
                assignList.put(((JAssignStmt) stmt).getLeftOp().toString(), ((JAssignStmt) stmt).getRightOp().toString());
                continue;
            }
            if (stmt instanceof JIfStmt) {
                String ifstms = ((JIfStmt) stmt).getCondition().toString();
                int nextUnitIndex = path.indexOf(stmt) + 1;
                Unit nextUnit = path.get(nextUnitIndex);

                //如果ifstmt的后继语句不是ifstmt中goto语句，说明ifstmt中的条件为假
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = "!( " + ifstms + " )";
                else
                    ifstms = "( " + ifstms + " )";
                stepConditionsWithJimpleVars.add(ifstms);
                continue;
            }
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
        System.out.println("The step conditions with JimpleVars are: " + stepConditionsWithJimpleVars);
        //stepConditionsWithJimpleVars:[!( $i1 != 0 )]
        //assignList:{$i1=i0 % 15}
        if (jVars.size() != 0) {
            //替换条件里的Jimple变量
            for (String cond : stepConditionsWithJimpleVars) {
                String tmp;
                while((tmp=containJvars(cond,jVars))!=null) {
                    cond=cond.replace(tmp, assignList.get(tmp).trim());
                }
                stepConditions.add(cond);
            }
            System.out.println("stepConditions:" + stepConditions);
            //stepConditions:[!( i0 % 15 != 0 )]

            for (Local lv : jVars) {
                if (expectedResult.contains(lv.toString())) {
                    expectedResult = expectedResult.replace(lv.toString(), assignList.get(lv.toString()).trim());
                }
            }
            //expectedResult:"FizzBuzz"
        } else {
            stepConditions = stepConditionsWithJimpleVars;
        }
        if (stepConditions.isEmpty())
            return "";
        pathConstraint = stepConditions.get(0);
        int i = 1;
        while (i < stepConditions.size()) {
            pathConstraint = pathConstraint + " && " + stepConditions.get(i);
            i++;
        }
        System.out.println("The path expression is: " + pathConstraint);
        //The path expression is: !( i0 % 15 != 0 )
        return pathConstraint;
    }

    public List<String> calExpectedResult(List<String> testSet,List<Local> locals) {
        List<String> resultSet = new ArrayList<>();
        try {
            for (String test : testSet) {
                String[] paras = test.split(" ");
                if(test.contains("null")){
                    resultSet.add("该路径无测试用例可达");
                    continue;
                }
                Object[] paras2=new Object[paras.length];
                Class<?>[] paraCls=new Class<?>[paras.length];
                for(int i=0;i<paras.length;i++){
                    String type=locals.get(i).getType().toString();
                    switch (type) {
                        case "int":
                            paraCls[i] = int.class;
                            paras2[i] = Integer.parseInt(paras[i].split("=")[1]);
                            break;
                        case "boolean":
                            paraCls[i] = boolean.class;
                            paras2[i] = Boolean.parseBoolean(paras[i].split("=")[1]);
                            break;
                        case "long":
                            paraCls[i] = long.class;
                            paras2[i] = Long.parseLong(paras[i].split("=")[1]);
                            break;
                        case "double":
                            paraCls[i] = double.class;
                            paras2[i] = Double.parseDouble(paras[i].split("=")[1]);
                            break;
                        case "short":
                            paraCls[i] = short.class;
                            paras2[i] = Short.parseShort(paras[i].split("=")[1]);
                            break;
                        case "char":
                            paraCls[i] = char.class;
                            paras2[i] = (paras[i].split("=")[1]).charAt(0);
                            break;
                        case "java.lang.Character":
                            paraCls[i] = Class.forName(type);
                            paras2[i] = (paras[i].split("=")[1]).charAt(0);
                            break;
                        case "java.lang.String":
                            paraCls[i] = Class.forName(type);
                            paras2[i] = paras[i].split("=")[1];
                            break;
                        default:
                            paraCls[i] = Class.forName(type);
                            paras2[i] = paras[i].split("=")[1];
                            break;
                    }
                }

                ExpectResult exp = new ExpectResult(clsPath, clsName, mtdName, paraCls);
                String res = exp.getResult(paras2).toString();
                resultSet.add(res);
            }
//            System.out.println(resultSet);
        }catch (Exception e){
            e.printStackTrace();
        }
        return resultSet;
    }

    public String solve(String pathConstraint) throws Exception {
        return Z3Solver.solve(pathConstraint);
    }

    public String randomTC(List<Local> parameters) {

        String varName;
        String varValue = "";
        String testinput = "";

        for (Local para : parameters) {
            varName = para.getName();
//            System.out.println("Type:"+para.getType().toString());
            if ("int".equals(para.getType().toString())) {
                varValue = String.valueOf((int) (Math.random() * 100));
            }
            else if ("boolean".equals(para.getType().toString())) {
                varValue = (int) (Math.random() * 100)%2==1?"true":"false";
            }
            else if ("long".equals(para.getType().toString())) {
                varValue = String.valueOf((long) (Math.random() * 100));
            }
            else if ("double".equals(para.getType().toString())) {
                varValue = String.valueOf((Math.random() * 100));
            }
            else if ("short".equals(para.getType().toString())) {
                varValue = String.valueOf((short)(Math.random() * 100));
            }
            else if ("char".equals(para.getType().toString())) {
                varValue = RandomStringUtils.randomAlphanumeric(10).substring(0,1);
            }
            else if ("java.lang.Character".equals(para.getType().toString())) {
                varValue = RandomStringUtils.randomAlphanumeric(10).substring(0,1);
            }
            else if ("java.lang.Integer".equals(para.getType().toString())) {
                varValue = String.valueOf((int) (Math.random() * 100));
            }
            else if ("java.lang.String".equals(para.getType().toString())) {
//                varValue = "abc";
                varValue = RandomStringUtils.randomAlphanumeric(10);
            }
            else if ("int[]".equals(para.getType().toString())) {
                int[] a = new int[5];
                for (int i = 0; i < 5; i++) {
                    a[i] = (int) (Math.random() * 100);
                }
                varValue = Arrays.toString(a);
            }
            testinput = testinput  + varName + "=" + varValue+" ";
//            System.out.println("input" + testinput);
        }
        return testinput;
    }

    public ArrayList<String> sortTestSet(List<String> testSet,List<Local> locals){
        ArrayList<String> sortedTest=new ArrayList<>();
        List<String> name=new ArrayList<>();
        locals.forEach(local -> name.add(local.getName()));
        for(String test:testSet){
            if(test.contains("null")){
                sortedTest.add(test);
            } else {
                HashMap<String, String> hashMap = new HashMap<>();
                String[] t = test.split(" ");
                for (String m : t) {
                    hashMap.put(m.split("=")[0], m.split("=")[1]);
                }
                StringBuilder res = new StringBuilder();
                for (String n : name) {
                    res.append(n).append("=").append(hashMap.get(n)).append(" ");
                }
                sortedTest.add(res.toString());
            }
        }
        return sortedTest;
    }

    public String containJvars(String cond,List<Local> jvars){
        for(Local local:jvars){
            if(cond.contains(local.toString())){
                return local.toString();
            }
        }
        return null;
    }

    public String solveString (Local local,String str,int round){
        if(round==0) {
            StringBuilder stringBuilder = new StringBuilder(str);
            String s = stringBuilder.substring(stringBuilder.indexOf("\"") + 1, stringBuilder.lastIndexOf("\""));
            return local.getName() + "=" + s;
        } else{
            return local.getName()+"="+RandomStringUtils.randomAlphanumeric(10);
        }
    }
}
